{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a52b664a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading from https://www.kaggle.com/api/v1/datasets/download/anshtanwar/pets-facial-expression-dataset?dataset_version_number=11...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 38.6M/38.6M [00:03<00:00, 12.0MB/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting files...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Path to dataset files: C:\\Users\\lenovo\\.cache\\kagglehub\\datasets\\anshtanwar\\pets-facial-expression-dataset\\versions\\11\n"
     ]
    }
   ],
   "source": [
    "import kagglehub\n",
    "\n",
    "# Download latest version\n",
    "path = kagglehub.dataset_download(\"anshtanwar/pets-facial-expression-dataset\")\n",
    "\n",
    "print(\"Path to dataset files:\", path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d081f5d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q kagglehub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5e4152d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "训练集类别： ['Angry', 'Other', 'Sad', 'happy']\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from torchvision import transforms, datasets\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# 数据集根路径（根据实际路径修改）\n",
    "master_folder = r\"C:\\Users\\lenovo\\.cache\\kagglehub\\datasets\\anshtanwar\\pets-facial-expression-dataset\\versions\\11\\Master Folder\"\n",
    "\n",
    "# 数据增强（训练集）\n",
    "train_transform = transforms.Compose([\n",
    "    transforms.RandomResizedCrop(224),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "# 验证集（无增强）\n",
    "val_transform = transforms.Compose([\n",
    "    transforms.Resize(256),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "# 加载数据集（绝对路径）\n",
    "train_dataset = datasets.ImageFolder(root=os.path.join(master_folder, \"train\"), transform=train_transform)\n",
    "val_dataset = datasets.ImageFolder(root=os.path.join(master_folder, \"valid\"), transform=val_transform)\n",
    "test_dataset = datasets.ImageFolder(root=os.path.join(master_folder, \"test\"), transform=val_transform)\n",
    "\n",
    "# 检查类别\n",
    "print(\"训练集类别：\", train_dataset.classes)\n",
    "\n",
    "# 定义数据加载器\n",
    "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)\n",
    "val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)\n",
    "test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "46ed005d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "训练集类别：['Angry', 'Other', 'Sad', 'happy']\n",
      "Epoch 1/20 | 训练 Loss: 1.2082 | 训练 Acc: 46.20% | 验证 Loss: 1.6943 | 验证 Acc: 38.89%\n",
      "→ 验证准确率提升！保存最佳模型至 cat_emotion_best_model.pth\n",
      "Epoch 2/20 | 训练 Loss: 0.8392 | 训练 Acc: 65.70% | 验证 Loss: 1.7169 | 验证 Acc: 47.22%\n",
      "→ 验证准确率提升！保存最佳模型至 cat_emotion_best_model.pth\n",
      "Epoch 3/20 | 训练 Loss: 0.6381 | 训练 Acc: 73.70% | 验证 Loss: 1.9550 | 验证 Acc: 44.44%\n",
      "Epoch 4/20 | 训练 Loss: 0.5921 | 训练 Acc: 77.60% | 验证 Loss: 1.7985 | 验证 Acc: 55.56%\n",
      "→ 验证准确率提升！保存最佳模型至 cat_emotion_best_model.pth\n",
      "Epoch 5/20 | 训练 Loss: 0.4407 | 训练 Acc: 84.00% | 验证 Loss: 1.3910 | 验证 Acc: 58.33%\n",
      "→ 验证准确率提升！保存最佳模型至 cat_emotion_best_model.pth\n",
      "Epoch 6/20 | 训练 Loss: 0.3039 | 训练 Acc: 88.90% | 验证 Loss: 2.0561 | 验证 Acc: 55.56%\n",
      "Epoch 7/20 | 训练 Loss: 0.2624 | 训练 Acc: 90.40% | 验证 Loss: 1.7313 | 验证 Acc: 52.78%\n",
      "Epoch 8/20 | 训练 Loss: 0.2264 | 训练 Acc: 92.70% | 验证 Loss: 1.8605 | 验证 Acc: 47.22%\n",
      "Epoch 9/20 | 训练 Loss: 0.1934 | 训练 Acc: 92.80% | 验证 Loss: 1.7900 | 验证 Acc: 52.78%\n",
      "Epoch 10/20 | 训练 Loss: 0.2432 | 训练 Acc: 90.30% | 验证 Loss: 2.3224 | 验证 Acc: 44.44%\n",
      "Epoch 11/20 | 训练 Loss: 0.2079 | 训练 Acc: 91.40% | 验证 Loss: 2.1658 | 验证 Acc: 47.22%\n",
      "Epoch 12/20 | 训练 Loss: 0.1951 | 训练 Acc: 93.00% | 验证 Loss: 2.3052 | 验证 Acc: 47.22%\n",
      "Epoch 13/20 | 训练 Loss: 0.1538 | 训练 Acc: 95.40% | 验证 Loss: 2.1611 | 验证 Acc: 41.67%\n",
      "Epoch 14/20 | 训练 Loss: 0.1598 | 训练 Acc: 94.80% | 验证 Loss: 1.9986 | 验证 Acc: 50.00%\n",
      "Epoch 15/20 | 训练 Loss: 0.1383 | 训练 Acc: 95.30% | 验证 Loss: 1.9466 | 验证 Acc: 50.00%\n",
      "Epoch 16/20 | 训练 Loss: 0.1234 | 训练 Acc: 95.40% | 验证 Loss: 1.9554 | 验证 Acc: 50.00%\n",
      "Epoch 17/20 | 训练 Loss: 0.1191 | 训练 Acc: 95.90% | 验证 Loss: 1.9468 | 验证 Acc: 41.67%\n",
      "Epoch 18/20 | 训练 Loss: 0.1117 | 训练 Acc: 95.50% | 验证 Loss: 2.0384 | 验证 Acc: 44.44%\n",
      "Epoch 19/20 | 训练 Loss: 0.1012 | 训练 Acc: 97.20% | 验证 Loss: 1.9767 | 验证 Acc: 41.67%\n",
      "Epoch 20/20 | 训练 Loss: 0.0987 | 训练 Acc: 96.90% | 验证 Loss: 1.9936 | 验证 Acc: 41.67%\n",
      "\n",
      "训练完成！最佳验证准确率：58.33% | 测试集准确率：50.00%\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import transforms, datasets, models\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# ====================== 配置参数（根据实际情况修改） ======================\n",
    "# 数据集路径（替换为你的Master Folder绝对路径）\n",
    "MASTER_FOLDER = r\"C:\\Users\\lenovo\\.cache\\kagglehub\\datasets\\anshtanwar\\pets-facial-expression-dataset\\versions\\11\\Master Folder\"\n",
    "TRAIN_DIR = os.path.join(MASTER_FOLDER, \"train\")  # 训练集路径\n",
    "VALID_DIR = os.path.join(MASTER_FOLDER, \"valid\")   # 验证集路径\n",
    "TEST_DIR = os.path.join(MASTER_FOLDER, \"test\")     # 测试集路径\n",
    "\n",
    "# 超参数\n",
    "BATCH_SIZE = 32       # 每批数据量（根据显存调整，如显存小设为16）\n",
    "NUM_EPOCHS = 20       # 训练轮数（小数据集20-30足够）\n",
    "LEARNING_RATE = 1e-3  # 学习率（默认Adam优化器的常用值）\n",
    "NUM_WORKERS = 0       # 数据加载线程数（Windows设为0，Linux设为4）\n",
    "\n",
    "# 设备（自动选择GPU或CPU）\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "\n",
    "# ====================== 数据加载与预处理 ======================\n",
    "# 数据增强（训练集：随机裁剪、翻转、标准化）\n",
    "train_transform = transforms.Compose([\n",
    "    transforms.RandomResizedCrop(224),          # 随机裁剪到224x224（模型输入尺寸）\n",
    "    transforms.RandomHorizontalFlip(p=0.5),     # 随机水平翻转（50%概率）\n",
    "    transforms.ToTensor(),                      # 转为Tensor（0-1归一化）\n",
    "    transforms.Normalize(                       # 标准化（预训练模型的统计量）\n",
    "        mean=[0.485, 0.456, 0.406], \n",
    "        std=[0.229, 0.224, 0.225]\n",
    "    )\n",
    "])\n",
    "\n",
    "# 验证/测试集预处理（无增强，仅resize和标准化）\n",
    "val_test_transform = transforms.Compose([\n",
    "    transforms.Resize(256),                     # 先resize到256x256\n",
    "    transforms.CenterCrop(224),                 # 中心裁剪到224x224\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(\n",
    "        mean=[0.485, 0.456, 0.406], \n",
    "        std=[0.229, 0.224, 0.225]\n",
    "    )\n",
    "])\n",
    "\n",
    "# 加载数据集（自动识别子文件夹为类别）\n",
    "try:\n",
    "    train_dataset = datasets.ImageFolder(root=TRAIN_DIR, transform=train_transform)\n",
    "    valid_dataset = datasets.ImageFolder(root=VALID_DIR, transform=val_test_transform)\n",
    "    test_dataset = datasets.ImageFolder(root=TEST_DIR, transform=val_test_transform)\n",
    "except FileNotFoundError as e:\n",
    "    raise ValueError(f\"路径错误：{e}。请检查MASTER_FOLDER是否正确指向包含train/valid/test的文件夹\")\n",
    "\n",
    "# 检查类别是否匹配（如['angry', 'happy', 'other', 'sad']）\n",
    "print(f\"训练集类别：{train_dataset.classes}\")\n",
    "NUM_CLASSES = len(train_dataset.classes)  # 自动获取类别数（无需手动修改）\n",
    "\n",
    "\n",
    "# ====================== 模型定义与初始化 ======================\n",
    "# 加载预训练模型（MobileNetV3 Small，轻量且适合小数据集）\n",
    "model = models.mobilenet_v3_small(pretrained=True)\n",
    "# 替换最后一层全连接层（输出为情绪类别数）\n",
    "model.classifier[-1] = nn.Linear(in_features=1024, out_features=NUM_CLASSES)\n",
    "model.to(DEVICE)  # 模型加载到GPU/CPU\n",
    "\n",
    "\n",
    "# ====================== 训练配置（损失函数、优化器、调度器） ======================\n",
    "criterion = nn.CrossEntropyLoss()  # 交叉熵损失（多分类任务）\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)  # Adam优化器\n",
    "# 学习率调度器（每5个epoch学习率乘以0.5，防止过拟合）\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)\n",
    "\n",
    "\n",
    "# ====================== 训练循环（含验证） ======================\n",
    "best_valid_acc = 0.0  # 记录最佳验证准确率\n",
    "save_path = \"cat_emotion_best_model.pth\"  # 最佳模型保存路径\n",
    "\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    # ---------------------- 训练阶段 ----------------------\n",
    "    model.train()  # 开启训练模式（启用Dropout/BatchNorm）\n",
    "    train_loss = 0.0\n",
    "    train_correct = 0\n",
    "    train_total = 0\n",
    "\n",
    "    for inputs, labels in DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS):\n",
    "        inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)  # 数据加载到GPU/CPU\n",
    "        \n",
    "        optimizer.zero_grad()  # 梯度清零\n",
    "        outputs = model(inputs)  # 前向传播（预测）\n",
    "        loss = criterion(outputs, labels)  # 计算损失\n",
    "        loss.backward()  # 反向传播（计算梯度）\n",
    "        optimizer.step()  # 优化器更新参数\n",
    "\n",
    "        # 统计训练指标\n",
    "        train_loss += loss.item() * inputs.size(0)  # 累计损失（乘以batch_size）\n",
    "        _, predicted = torch.max(outputs, 1)  # 获取预测类别（最大值索引）\n",
    "        train_correct += (predicted == labels).sum().item()  # 累计正确数\n",
    "        train_total += labels.size(0)  # 累计总样本数\n",
    "\n",
    "    # 计算训练指标\n",
    "    train_loss_epoch = train_loss / train_total\n",
    "    train_acc_epoch = 100.0 * train_correct / train_total\n",
    "\n",
    "    # ---------------------- 验证阶段 ----------------------\n",
    "    model.eval()  # 开启评估模式（禁用Dropout/BatchNorm）\n",
    "    valid_loss = 0.0\n",
    "    valid_correct = 0\n",
    "    valid_total = 0\n",
    "\n",
    "    with torch.no_grad():  # 不计算梯度（节省显存）\n",
    "        for inputs, labels in DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS):\n",
    "            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)\n",
    "            \n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "            \n",
    "            valid_loss += loss.item() * inputs.size(0)\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            valid_correct += (predicted == labels).sum().item()\n",
    "            valid_total += labels.size(0)\n",
    "\n",
    "    # 计算验证指标\n",
    "    valid_loss_epoch = valid_loss / valid_total\n",
    "    valid_acc_epoch = 100.0 * valid_correct / valid_total\n",
    "\n",
    "    # 更新学习率调度器\n",
    "    scheduler.step()\n",
    "\n",
    "    # 打印训练日志\n",
    "    print(f\"Epoch {epoch+1}/{NUM_EPOCHS} | \"\n",
    "          f\"训练 Loss: {train_loss_epoch:.4f} | 训练 Acc: {train_acc_epoch:.2f}% | \"\n",
    "          f\"验证 Loss: {valid_loss_epoch:.4f} | 验证 Acc: {valid_acc_epoch:.2f}%\")\n",
    "\n",
    "    # 保存最佳验证准确率的模型\n",
    "    if valid_acc_epoch > best_valid_acc:\n",
    "        best_valid_acc = valid_acc_epoch\n",
    "        torch.save(model.state_dict(), save_path)\n",
    "        print(f\"→ 验证准确率提升！保存最佳模型至 {save_path}\")\n",
    "\n",
    "\n",
    "# ====================== 最终测试（使用测试集） ======================\n",
    "model.load_state_dict(torch.load(save_path))  # 加载最佳模型\n",
    "model.eval()\n",
    "\n",
    "test_correct = 0\n",
    "test_total = 0\n",
    "with torch.no_grad():\n",
    "    for inputs, labels in DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS):\n",
    "        inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)\n",
    "        outputs = model(inputs)\n",
    "        _, predicted = torch.max(outputs, 1)\n",
    "        test_correct += (predicted == labels).sum().item()\n",
    "        test_total += labels.size(0)\n",
    "\n",
    "test_acc = 100.0 * test_correct / test_total\n",
    "print(f\"\\n训练完成！最佳验证准确率：{best_valid_acc:.2f}% | 测试集准确率：{test_acc:.2f}%\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Pytorch_cpu",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
